import numpy as np
import torch
import pandas as pd
from torch.utils.data import DataLoader, Dataset, TensorDataset
from tumpreprocess import GtInterpolate
from tumpreprocess import LoadTUMData

def main():

    res_train, res_val = readTUM()
    for i, (batch_0, batch_1) in enumerate(res_train):
        print(i, batch_1)
        break
    return

class TUMdataset(Dataset):
    def __init__(self, group):
        self.position = group[1]
        self.gt_position = group[0]

    def __len__(self):
        return len(self.position)

    def __getitem__(self, idx):
        # position = self.data_frame[idx][1]
        #### not sensor data but position, need integration
        # gt_position = self.data_frame[idx][0]
        return self.position[:][idx], self.gt_position[:][idx]


def readTUM():
    # train_data = TUMdataset('/home/ubuntu/user_space/TUM/.....')

    time, tum_accex, tum_accey, tum_accez, tum_angux, tum_anguy, tum_anguz = LoadTUMData.loadTumImu()
    tum_px, tum_py, tum_pz, tum_q0, tum_q1, tum_q2, tum_q3 = LoadTUMData.loadTumGT()
    group = GtInterpolate.getRandomGroup(time, tum_px, tum_py, tum_pz, tum_accex, tum_accey, tum_accez, tum_angux, tum_anguy, tum_anguz)

    ### zheme jiandan de wenti douyou bug, dou gaobudong dou huicuo
    position = []
    for i in range(len(group[13])):
        position.append([group[7][i], group[8][i], group[9][i], group[10][i], group[11][i], group[12][i], 0])

    train_group = [group[13], position]
    train_data = TUMdataset(train_group)

    print(type(train_data))



    train_pos = torch.from_numpy(np.array(position))
    train_gt = torch.from_numpy(np.array(group[13]))

    train_pos = train_pos.float()
    train_gt = train_gt.float()
    train_dataset = TensorDataset(train_pos, train_gt)
    # train_dataset = TensorDataset(train_pos)
    val_dataset = TensorDataset(train_gt)
    train_dataloader = DataLoader(train_dataset, batch_size=7, shuffle=False, drop_last=True)
    val_dataloader = DataLoader(val_dataset, batch_size=7, shuffle=False, drop_last=True)

    # train_dataloader = DataLoader(train_data, batch_size=7, shuffle=False)

    return train_dataloader, val_dataloader






if __name__ == '__main__':
     main()
